import pandas as pd
import numpy as np
from datetime import datetime, timedelta
from typing import List, Dict, Any
import concurrent.futures
from tqdm import tqdm

from data.market_data_fetcher import MarketDataFetcher


class MAStrategy:
    def __init__(self, max_workers: int = 10):
        self.market_data = MarketDataFetcher()
        self.max_workers = max_workers

    def calculate_ma(self, df: pd.DataFrame, periods: List[int]) -> pd.DataFrame:
        """Calculate moving averages for given periods"""
        for period in periods:
            df[f'MA{period}'] = df['close'].rolling(window=period).mean()
        return df

    def is_ma_uptrend(self, ma_values: pd.Series, days: int = 20) -> bool:
        """Check if MA is in uptrend for the last n days"""
        if len(ma_values) < days:
            return False
        recent_ma = ma_values[-days:].values  # 转换为numpy数组
        return all(recent_ma[i] >= recent_ma[i-1] for i in range(1, len(recent_ma)))

    def analyze_single_stock(self, stock_code: str, market_info: Dict[str, Any]) -> Dict[str, Any]:
        """分析单个股票是否符合策略条件"""
        try:
            # 获取历史数据
            df = self.market_data.get_stock_data(stock_code)
            
            # 过滤数据记录少于600条的股票
            if len(df) < 600:
                return None
                
            # 计算移动平均线
            df = self.calculate_ma(df, [20, 60, 100])
            
            # 获取最新数据点
            latest = df.iloc[-1]
            
            # 检查条件
            # 3. 股价高于60日和100日均线
            if not (latest['close'] > latest['MA60'] and latest['close'] > latest['MA100']):
                return None
                
            # 4. 60日均线高于100日均线
            if not (latest['MA60'] > latest['MA100']):
                return None
                
            # 5. 100日均线至少上涨20天
            if not self.is_ma_uptrend(df['MA100'], 20):
                return None
                
            # 6. 20日均线高于60日和100日均线
            if not (latest['MA20'] > latest['MA60'] and latest['MA20'] > latest['MA100']):
                return None
                
            # 7. 目前的股价比100日内最低点至少高出25%
            min_price_100d = df['close'][-100:].min()
            if not (latest['close'] >= min_price_100d * 1.25):
                return None
                
            # 8. 目前的股价处于100日最高点的25%以内
            max_price_100d = df['close'][-100:].max()
            if not (latest['close'] >= max_price_100d * 0.75):
                return None
            
            # 返回市场信息
            return market_info
            
        except Exception as e:
            print(f"分析股票 {stock_code} 时发生错误: {str(e)}")
            return None

    def run_strategy(self) -> pd.DataFrame:
        """
        执行均线选股策略：
        1. 市值大于100亿
        2. 数据记录大于600条
        3. 股价高于60日和100日均线
        4. 60日均线高于100日均线
        5. 100日均线至少上涨20天
        6. 20日均线高于60日和100日均线
        7. 目前的股价比100日内最低点至少高出25%
        8. 目前的股价处于100日最高点的25%以内
        """
        print("开始获取实时市场数据...")
        # 获取实时市场数据
        market_data = self.market_data.get_realtime_market_data()
        
        # 过滤市值大于100亿的股票
        market_data['market_value'] = market_data['total_market_cap'].astype(float)
        filtered_stocks = market_data[market_data['market_value'] > 10000000000]
        print(f"找到 {len(filtered_stocks)} 只市值大于100亿的股票")

        # 创建股票代码到市场信息的映射
        stock_info_map = {
            row['ts_code']: {
                'ts_code': row['ts_code'],
                'name': row['name'],
                'close': row['close'],
                'change_percent': row['pct_change'],
                'volume': row['volume'],
                'market_value': row['market_value']
            }
            for _, row in filtered_stocks.iterrows()
        }

        selected_stocks = []
        
        print(f"使用 {self.max_workers} 个线程进行分析...")
        # 使用线程池处理股票分析
        with concurrent.futures.ThreadPoolExecutor(max_workers=self.max_workers) as executor:
            # 创建future到股票代码的映射
            future_to_stock = {
                executor.submit(self.analyze_single_stock, stock_code, stock_info): stock_code
                for stock_code, stock_info in stock_info_map.items()
            }
            
            # 使用tqdm显示进度
            for future in tqdm(concurrent.futures.as_completed(future_to_stock),
                             total=len(future_to_stock),
                             desc="分析股票进度"):
                stock_code = future_to_stock[future]
                try:
                    result = future.result()
                    if result is not None:
                        selected_stocks.append(result)
                except Exception as e:
                    print(f"处理股票 {stock_code} 时发生错误: {str(e)}")
        
        # 转换结果为DataFrame
        if selected_stocks:
            result_df = pd.DataFrame(selected_stocks)
            result_df = result_df.sort_values('market_value', ascending=False)
            print(f"\n找到 {len(result_df)} 只符合条件的股票")
            # 获取当前日期
            current_date = datetime.now().strftime("%Y%m%d")  # 格式化为 'YYYYMMDD'

            # 拼接文件名
            file_name = f"stocks_{current_date}_selected.csv"

            # 写入到 CSV 文件
            result_df.to_csv(file_name, index=False)

            print(f"数据已成功写入到文件：{file_name}")

            return result_df
        
        print("\n没有找到符合条件的股票")
        return pd.DataFrame()

    def get_strategy_results(self) -> pd.DataFrame:
        """获取策略结果"""
        return self.run_strategy()
